[Feature] Optimize forward metadata collection across dp ranks #1593
Conversation
|
@NeverRaR PTAL |
|
lgtm |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
There was a problem hiding this comment.
Pull Request Overview
This PR optimizes how forward-pass metadata is collected and communicated across data-parallel ranks by removing the previous all-reduce and introducing an HCCL-based all-gather approach.
- Enforce that dummy batch execution only runs under data parallelism and refactor
execute_dummy_batchto use per-rank metadata. - Replace
dist.all_reducewith HCCLall_gatherin_get_forward_metadata_across_dpand update callers to handle the Tensor of per-rank token counts. - Propagate
num_tokens_across_dpthrough dummy runs and forward contexts, masking sentinel values before the pass.
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| vllm_ascend/worker/worker_v1.py | Added assertion for dp_size > 1, refactored dummy-run logic to use HCCL per-rank metadata. |
| vllm_ascend/worker/model_runner_v1.py | Swapped all_reduce for all_gather under get_dp_group(), changed method signature and updated callers to handle a Tensor of metadata. |
Comments suppressed due to low confidence (1)
vllm_ascend/worker/model_runner_v1.py:622
- Add unit or integration tests covering the
dp_size > 1aggregation path to verify thatall_gatherproduces the correct combined metadata and that themasked_fill_logic correctly replaces sentinel values.
local_forward_metadata)
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #1593 +/- ##
===========================================
+ Coverage 27.39% 53.92% +26.52%
===========================================
Files 56 80 +24
Lines 6191 9964 +3773
===========================================
+ Hits 1696 5373 +3677
- Misses 4495 4591 +96
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
jianzs
added
performance-test
|
@Yikun @wangxiyuan @ApsarasX @ganyi1996ppo ready to merge. |
| max_num_tokens) | ||
| runner._dummy_run(max_num_tokens, | ||
| else: | ||
| num_tokens = 1 |
There was a problem hiding this comment.
If graph mode is off, a dummy run only needs to be executed; computational requirements are not a factor.
jianzs
force-pushed
the
feat/dp-comm-opt
branch
3 times, most recently
from
bfec114 to
dbaad95
Compare
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
…rker Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
Co-authored-by: Angazenn <92204292+Angazenn@users.noreply.github.com> Signed-off-by: Jade Zheng <zheng.shoujian@outlook.com>
|
@Angazenn PTAL |
|
@wangxiyuan @ganyi1996ppo @Yikun ready to merge. |
What this PR does / why we need it?
This PR introduces two optimizations for cases where data parallel size > 1:
set_forward_contextDoes this PR introduce any user-facing change?
no
How was this patch tested?
CI passed.